package com.javaee.wordtree.controller;

import com.javaee.wordtree.entity.DeepseekResponse;
import com.javaee.wordtree.entity.Word;
import com.javaee.wordtree.service.DeepseekService;
import com.javaee.wordtree.service.MemoryService;

import lombok.Data;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import javax.annotation.Resource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.RestTemplate;
import reactor.core.publisher.Mono;

@RestController
@RequestMapping("/api/deepseek")
public class DeepseekController {

    private static final Logger log = LoggerFactory.getLogger(DeepseekController.class);

    @Resource
    MemoryService memoryService;

    private final DeepseekService deepseekService;
    private final RestTemplate restTemplate;

    // 图像生成服务URL
    @Value("${image.generation.url:http://127.0.0.1:8080/generate}")
    private String imageGenerationUrl;

    // 存储用户故事状态
    private final Map<String, StoryState> userStoryStates = new ConcurrentHashMap<>();

    public DeepseekController(DeepseekService deepseekService) {
        this.deepseekService = deepseekService;
        this.restTemplate = new RestTemplate();
    }

    @PostMapping("ask")
    public Mono<ResponseEntity<DeepseekResponse>> getChatCompletion(
            @RequestBody ChatRequest chatRequest) {
        return deepseekService.callDeepseekApi(chatRequest.getQuestion())
                .map(ResponseEntity::ok)
                .onErrorResume(e -> Mono.just(
                        ResponseEntity.internalServerError()
                                .body(new DeepseekResponse())));
    }

    @Data
    public static class ChatRequest {
        private String question;
    }

    @PostMapping("story")
    public Mono<ResponseEntity<DeepseekResponse>> generateStory(@RequestParam String userID) {
        log.info("开始生成互动故事 - 用户ID: {}", userID);
        List<Word> learnedWords = memoryService.getTodayLearnedWords(userID);
        List<String> wordList = learnedWords.stream()
                .map(Word::getWord)
                .collect(Collectors.toList());

        if (wordList.isEmpty()) {
            // 如果没有学习单词，使用一些基础词汇
            wordList = List.of("adventure", "journey", "discover", "mystery", "explore");
            log.info("用户未学习单词，使用默认单词列表");
        }
        final List<String> finalWordList = wordList;

        log.info("使用单词列表生成故事: {}", String.join(", ", wordList));
        String prompt = buildIntroPrompt(wordList);

        // 创建并保存故事状态
        StoryState storyState = new StoryState();
        storyState.words = wordList;
        storyState.unusedWords = wordList;
        userStoryStates.put(userID, storyState);

        return deepseekService.callDeepseekApi(prompt)
                .map(response -> {
                    if (response != null && response.getChoices() != null &&
                        !response.getChoices().isEmpty() &&
                        response.getChoices().get(0).getMessage() != null) {

                        String content = response.getChoices().get(0).getMessage().getContent();

                        // 更新故事状态
                        storyState.storyHistory = content;

                        // 检查哪些单词已经使用
                        updateUnusedWords(storyState);

                        // 生成故事场景图像
                        try {
                            String sceneImage = generateSceneImage(content);
                            storyState.currentSceneImage = sceneImage;
                            response.setSceneImage(sceneImage);
                        } catch (Exception e) {
                            log.error("生成故事场景图像失败", e);
                        }
                    }
                    // 在生成 DeepseekResponse 后，添加关键词数组
                    response.setKeywords(finalWordList);
                    return ResponseEntity.ok(response);
                })
                .onErrorResume(e -> {
                    log.error("生成故事失败", e);
                    return Mono.just(
                            ResponseEntity.internalServerError()
                                    .body(new DeepseekResponse()));
                });
    }

    @PostMapping("story/choice")
    public Mono<ResponseEntity<DeepseekResponse>> processChoice(
            @RequestParam String userID,
            @RequestParam int choiceIndex) {

        log.info("处理用户故事选择 - 用户ID: {}, 选择: {}", userID, choiceIndex);

        // 获取用户的故事状态
        StoryState storyState = userStoryStates.get(userID);
        if (storyState == null) {
            log.error("未找到用户故事状态");
            return Mono.just(
                    ResponseEntity.badRequest()
                            .body(new DeepseekResponse()));
        }

        // 添加用户选择到历史记录
        storyState.userChoices.add(choiceIndex);
        storyState.branchCount = storyState.branchCount + 1;

        // 决定是否应该结束故事
        boolean shouldEndStory = storyState.userChoices.size() >= 4 ||
                               storyState.unusedWords.isEmpty();

        String prompt;
        if (shouldEndStory) {
            // 生成故事结局
            prompt = buildEndingPrompt(storyState.storyHistory, storyState.words);
        } else {
            // 生成故事分支
            prompt = buildBranchPrompt(storyState.storyHistory,
                                       storyState.words,
                                       storyState.unusedWords);
        }

        return deepseekService.callDeepseekApi(prompt)
                .map(response -> {
                    if (response != null && response.getChoices() != null &&
                        !response.getChoices().isEmpty() &&
                        response.getChoices().get(0).getMessage() != null) {

                        String content = response.getChoices().get(0).getMessage().getContent();

                        // 更新故事历史
                        storyState.storyHistory = storyState.storyHistory + "\n\n" + content;

                        // 更新未使用的单词
                        updateUnusedWords(storyState);

                        // 生成新的故事场景图像
                        try {
                            String sceneImage = generateSceneImage(content);
                            storyState.currentSceneImage = sceneImage;
                            response.setSceneImage(sceneImage);
                        } catch (Exception e) {
                            log.error("生成故事场景图像失败", e);
                        }

                        // 如果是结局，则可以清理状态
                        if (shouldEndStory) {
                            // 保留一段时间后清理或立即清理
                            // userStoryStates.remove(userID);
                        }
                    }
                    return ResponseEntity.ok(response);
                })
                .onErrorResume(e -> {
                    log.error("处理故事选择失败", e);
                    return Mono.just(
                            ResponseEntity.internalServerError()
                                    .body(new DeepseekResponse()));
                });
    }

    // 从故事文本中提取图像描述并生成图像
    private String generateSceneImage(String storyText) {
        try {
            // 1. 提取故事中适合生成图像的描述
            String imagePrompt = extractImagePrompt(storyText);

            // 2. 调用图像生成服务
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);

            Map<String, Object> requestBody = new ConcurrentHashMap<>();
            requestBody.put("prompts", List.of(imagePrompt));
            requestBody.put("num_inference_steps", 30);  // 可调整步数
            requestBody.put("guidance_scale", 7.5);      // 可调整指导比例

            HttpEntity<Map<String, Object>> request = new HttpEntity<>(requestBody, headers);

            log.info("正在生成场景图像，提示词: {}", imagePrompt);
            ResponseEntity<Map> response = restTemplate.postForEntity(
            imageGenerationUrl, 
            request, 
            Map.class
);
            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                @SuppressWarnings("unchecked")
                List<String> images = (List<String>) response.getBody().get("generated_images");
                if (images != null && !images.isEmpty()) {
                    log.info("场景图像生成成功");
                    return images.get(0);  // 返回第一张图像的base64编码
                }
            }

            log.error("图像生成服务返回异常: {}", response);
            return null;
        } catch (Exception e) {
            log.error("调用图像生成服务失败: URL={}, 请求体={}, 错误信息={}",
                    imageGenerationUrl,  e.getMessage(), e);
            return null;
        }
    }

    // 从故事文本中提取适合生成图像的描述
    private String extractImagePrompt(String storyText) {
        // 构建提取图像描述的提示
        String extractionPrompt = "You are a professional AI prompt writer for text-to-image models. "
                + "Given an English story fragment, extract a concise, vivid description for a single-scene illustration. "
                + "Output ONLY the prompt without any explanations or formatting. "
                + "Focus on the main character and setting with descriptive details. Story fragment: " + storyText;

        try {
            DeepseekResponse response = deepseekService.callDeepseekApi(extractionPrompt).block();
            if (response != null && response.getChoices() != null && !response.getChoices().isEmpty()
                    && response.getChoices().get(0).getMessage() != null) {
                String imagePrompt = response.getChoices().get(0).getMessage().getContent();
                // 添加风格指导
                imagePrompt += ", 8-bit, simple, colorful, single scene, cute, clear composition, centered subject, clean background";
                return imagePrompt;
            }
        } catch (Exception e) {
            log.error("提取图像描述失败: {}", e.getMessage(), e);
        }

        // 如果提取失败，使用故事的前100个字符加上风格描述
        String fallbackPrompt = storyText;
        if (fallbackPrompt.length() > 100) {
            fallbackPrompt = fallbackPrompt.substring(0, 100);
        }
        return fallbackPrompt + ", 8-bit, simple, colorful, single scene, cute, clear composition, centered subject, clean background";
    }

    // 辅助方法：更新未使用的单词列表
    private void updateUnusedWords(StoryState storyState) {
        String lowerCaseStory = storyState.storyHistory.toLowerCase();
        List<String> unusedWords = storyState.words.stream()
                .filter(word -> !lowerCaseStory.contains(word.toLowerCase()))
                .collect(Collectors.toList());
        storyState.unusedWords = unusedWords;
    }

    // 辅助方法：构建初始故事提示
private String buildIntroPrompt(List<String> words) {
    return "You are an advanced English learning story AI. "
        + "Generate a creative and engaging story introduction that uses ALL the following words at least once: "
        + String.join(", ", words) + ". "
        + "Highlight the target words by enclosing them in asterisks (*). "
        + "First write the story in English, then provide a Chinese translation. "
        + "At the end, provide 3 numbered choices for the next action (choices only in English). "
        + "Format:\n\n"
        + "English:\n<Story introduction in English>\n\n"
        + "Chinese:\n<故事介绍的中文翻译>\n\n"
        + "Choices:\n1. <Option 1>\n2. <Option 2>\n3. <Option 3>";
}

// 修改故事分支提示
private String buildBranchPrompt(String storyHistory, List<String> words, List<String> unusedWords) {
    return "Continue the interactive English story using the following context and user choices.\n" +
           "Story so far:\n" +
           storyHistory + "\n" +
           "Remaining target words: " + String.join(", ", unusedWords) + " (if any not used yet)\n\n" +
           "Write the next part of the story first in English, then in Chinese. At the end, provide 2 or 3 new choices (in English only).\n\n" +
           "English:\n" +
           "<Next part of the story in English>\n\n" +
           "Chinese:\n" +
           "<故事下一部分的中文翻译>\n\n" +
           "Choices:\n" +
           "1. <Option 1>\n" +
           "2. <Option 2>";
}

// 修改故事结局提示
private String buildEndingPrompt(String storyHistory, List<String> words) {
    return "Based on the story context and user choices below, write a logical and satisfying ending for the interactive English story.\n" +
           "Make sure all target words " + String.join(", ", words) + " have been used at least once in the story.\n\n" +
           "Story so far:\n" +
           storyHistory + "\n\n" +
           "Write the story ending first in English, then in Chinese.\n\n" +
           "English:\n" +
           "<Ending text in English>\n\n" +
           "Chinese:\n" +
           "<结尾文本的中文翻译>";
}

    // 故事状态类 - 添加图像字段
    public static class StoryState {
        public List<String> words;
        public List<String> unusedWords;
        public String storyHistory = "";
        public List<Integer> userChoices = new ArrayList<>();
        public int branchCount = 1;
        public String currentSceneImage; // 存储当前场景图像的base64编码
    }
}